from pymilvus import (
    connections,
    Collection,
)
from zhipuai import ZhipuAI
import os
from dotenv import load_dotenv, find_dotenv


def insert_texts_to_milvus(texts_to_insert):
    """
    将一组文本转换为向量并插入Milvus集合中。
    
    参数:
    texts_to_insert (List[str]): 要插入Milvus的文本列表。
    """
    # 加载环境变量
    load_dotenv(find_dotenv())
    api_key = os.getenv('ZHIPU_API_KEY')

    # 初始化ZhipuAI客户端
    client = ZhipuAI(api_key=api_key)

    # 确保连接Milvus
    connections.connect("default", host='127.0.0.1', port='19530')

    # 获取集合实例
    collection = Collection("text_collection")

    # 准备数据并插入
    for text in texts_to_insert:
        # 将文本转换为向量
        response = client.embeddings.create(
            model="embedding-2",
            input=text,
        )
        vector = response.data[0].embedding

        # 准备数据，假设"id"会自动生成，所以我们只提供"text"和"vector"
        data = [
            {"text": text, "vector": vector}
        ]

        # 插入数据
        mr = collection.insert(data)
        collection.flush()  # 刷新以确保数据被写入

        # 如果需要，可以获取插入的实体ID
        ids = mr.primary_keys
        print(f"Inserted text '{text}' with ID: {ids} into Milvus.")

    print("All texts have been inserted.")


if __name__ == "__main__":
    # 示例用法
    texts_to_insert_example = ["你好"]
    insert_texts_to_milvus(texts_to_insert_example)
